import os
import requests
import json
from fastmcp import FastMCP, Client
import asyncio

API_BASE_URL = "https://www.alphavantage.co/query"

mcp = FastMCP('demo.mcp')

def _make_api_request(function_name: str, symbol: str) -> str:
    """
    Helper function to make API requests to Alpha Vantage and handle responses.
    
    Input parameters:
    function_name: The Alpha Vantage API function to call (e.g., 'OVERVIEW').
    symbol: The stock symbol for the API request (e.g., 'AAPL').
    
    Output parameters:
    result: Returns a JSON string containing the API response or an error message.
    """
    try:
        api_key = os.getenv("ALPHA_VANTAGE_API_KEY")
        if not api_key:
            raise ValueError("ALPHA_VANTAGE_API_KEY environment variable is not set.")
        
        params = {
            "function": function_name,
            "symbol": symbol,
            "apikey": api_key,
            "source": "trading_agents",
        }
        
        response = requests.get(API_BASE_URL, params=params)
        response.raise_for_status()
        
        response_json = response.json()
        if "Information" in response_json:
            info_message = response_json["Information"]
            if "rate limit" in info_message.lower() or "api key" in info_message.lower():
                return json.dumps({"error": f"Alpha Vantage rate limit exceeded: {info_message}"})
        
        return json.dumps(response_json)
    except Exception as e:
        return json.dumps({"error": str(e)})

@mcp.tool()
def get_company_overview(symbol: str) -> str:
    """
    Retrieves company overview information from Alpha Vantage API for a given stock symbol.
    
    Input parameters:
    symbol: The stock symbol for which to retrieve company overview data (e.g., 'AAPL' for Apple Inc.). This parameter is required and should be a string.
    
    Output parameters:
    result: Returns a JSON string containing the company overview data or an error message.
    """
    return _make_api_request("OVERVIEW", symbol)

@mcp.tool()
def get_balance_sheet(symbol: str) -> str:
    """
    Retrieves balance sheet information from Alpha Vantage API for a given stock symbol.
    
    Input parameters:
    symbol: The stock symbol for which to retrieve balance sheet data (e.g., 'AAPL' for Apple Inc.). This parameter is required and should be a string.
    
    Output parameters:
    result: Returns a JSON string containing the balance sheet data or an error message.
    """
    return _make_api_request("BALANCE_SHEET", symbol)

@mcp.tool()
def get_cash_flow(symbol: str) -> str:
    """
    Retrieves cash flow information from Alpha Vantage API for a given stock symbol.
    
    Input parameters:
    symbol: The stock symbol for which to retrieve cash flow data (e.g., 'AAPL' for Apple Inc.). This parameter is required and should be a string.
    
    Output parameters:
    result: Returns a JSON string containing the cash flow data or an error message.
    """
    return _make_api_request("CASH_FLOW", symbol)

@mcp.tool()
def get_income_statement(symbol: str) -> str:
    """
    Retrieves income statement information from Alpha Vantage API for a given stock symbol.
    
    Input parameters:
    symbol: The stock symbol for which to retrieve income statement data (e.g., 'AAPL' for Apple Inc.). This parameter is required and should be a string.
    
    Output parameters:
    result: Returns a JSON string containing the income statement data or an error message.
    """
    return _make_api_request("INCOME_STATEMENT", symbol)

ALPHA_VANTAGE_COMPANY_OVERVIEW = {
    "type": "function",
    "function": {
        "name": "get_company_overview",
        "description": "Retrieves company overview information from Alpha Vantage API for a given stock symbol.",
        "parameters": {
            "type": "object",
            "properties": {
                "symbol": {
                    "type": "string",
                    "description": "The stock symbol for which to retrieve company overview data (e.g., 'AAPL' for Apple Inc.)."
                }
            },
            "required": ["symbol"]
        }
    }
}

ALPHA_VANTAGE_BALANCE_SHEET = {
    "type": "function",
    "function": {
        "name": "get_balance_sheet",
        "description": "Retrieves balance sheet information from Alpha Vantage API for a given stock symbol.",
        "parameters": {
            "type": "object",
            "properties": {
                "symbol": {
                    "type": "string",
                    "description": "The stock symbol for which to retrieve balance sheet data (e.g., 'AAPL' for Apple Inc.)."
                }
            },
            "required": ["symbol"]
        }
    }
}

ALPHA_VANTAGE_CASH_FLOW = {
    "type": "function",
    "function": {
        "name": "get_cash_flow",
        "description": "Retrieves cash flow information from Alpha Vantage API for a given stock symbol.",
        "parameters": {
            "type": "object",
            "properties": {
                "symbol": {
                    "type": "string",
                    "description": "The stock symbol for which to retrieve cash flow data (e.g., 'AAPL' for Apple Inc.)."
                }
            },
            "required": ["symbol"]
        }
    }
}

ALPHA_VANTAGE_INCOME_STATEMENT = {
    "type": "function",
    "function": {
        "name": "get_income_statement",
        "description": "Retrieves income statement information from Alpha Vantage API for a given stock symbol.",
        "parameters": {
            "type": "object",
            "properties": {
                "symbol": {
                    "type": "string",
                    "description": "The stock symbol for which to retrieve income statement data (e.g., 'AAPL' for Apple Inc.)."
                }
            },
            "required": ["symbol"]
        }
    }
}

async def main():
    client = Client(mcp)
    async with client:
        # View available tools
        tools = await client.list_tools()
        print('Available tools:', tools)
        
        # Call the tools (parameter names must match)
        overview_result = await client.call_tool('get_company_overview', {'symbol': 'AAPL'})
        print('Company Overview Result:', overview_result)
        
        balance_sheet_result = await client.call_tool('get_balance_sheet', {'symbol': 'AAPL'})
        print('Balance Sheet Result:', balance_sheet_result)
        
        cash_flow_result = await client.call_tool('get_cash_flow', {'symbol': 'AAPL'})
        print('Cash Flow Result:', cash_flow_result)
        
        income_statement_result = await client.call_tool('get_income_statement', {'symbol': 'AAPL'})
        print('Income Statement Result:', income_statement_result)

if __name__ == '__main__':
    asyncio.run(main())